package com.qy.tree;

public class RBTreeDemo<K extends Comparable<K>,V> {

    private static final Boolean RED = true;
    private static final Boolean BLACK = false;
    private RBNode<K,V> root;

    public Boolean colorOf(RBNode<K,V> node){
        return node == null ? BLACK : node.color;
    }
    public void setColor(RBNode<K,V> node,Boolean color){
        if(node != null){
            node.color = color;
        }
    }
    public RBNode<K,V> leftOf(RBNode<K,V> node){
        return node == null ? null : node.left;
    }
    public RBNode<K,V> rightOf(RBNode<K,V> node){
        return node == null ? null : node.right;
    }
    public RBNode<K,V> parentOf(RBNode<K,V> node){
        return node == null ? null : node.parent;
    }

    public void leftRotate(RBNode<K,V> p){
        if(p != null){
            RBNode<K,V> r = p.right;
            p.right = r.left;
            if(r.left != null){
                r.left.parent = p;
            }
            r.parent = p.parent;
            if(p.parent == null){
                root = r;
            }else if(p.parent.left == p){
                p.parent.left = r;
            }else{
                p.parent.right = r;
            }
            r.left = p;
            p.parent = r;
        }
    }

    public void rightRotate(RBNode<K,V> p){
        if(p != null){
            RBNode<K,V> l = p.left;
            p.left = l.right;
            if(l.right != null){
                l.right.parent = p;
            }
            l.parent = p.parent;
            if(p.parent == null){
                root = l;
            }else if(p.parent.left == p){
                p.parent.left = l;
            }else{
                p.parent.right = l;
            }
            l.right = p;
            p.parent = l;
        }
    }

    public void put(K key,V value){
        RBNode<K,V> temp = root;
        if(temp == null){
            root = new RBNode<K,V>(key,value,null);
            return;
        }
        int cmp;
        RBNode<K,V> parent;
        do{
            cmp = key.compareTo(temp.key);
            parent = temp;
            if(cmp < 0){
                temp = temp.left;
            }else if(cmp > 0){
                temp = temp.right;
            }else{
                temp.value = value;
                return;
            }
        }while(temp != null);
        RBNode<K,V> newNode = new RBNode<K,V>(key,value,parent);
        if(cmp < 0){
            parent.left = newNode;
        }else{
            parent.right = newNode;
        }
        fixAfterPut(newNode);

    }

    private void fixAfterPut(RBNode<K,V> newNode) {
        //插入就是红色的
        newNode.color = RED;
        //2节点无需调整，3节点和4节点需要调整
        while(newNode != null && newNode != root && colorOf(parentOf(newNode)) == RED){
            //分两种父亲在爷爷左或右，因为是3,4节点
            if(parentOf(newNode) == leftOf(parentOf(parentOf(newNode)))){
                //看是不是4节点，有没有叔叔
                RBNode<K,V> y = rightOf(parentOf(parentOf(newNode)));
                if(colorOf(y) == RED){
                    //有叔叔
                    setColor(parentOf(parentOf(newNode)),RED);
                    setColor(parentOf(newNode),BLACK);
                    setColor(y,BLACK);
                    newNode = parentOf(parentOf(newNode));//递归到为空或是根节点
                }else{
                    //没有叔叔
                    if(newNode == rightOf(parentOf(newNode))){
                        newNode = parentOf(newNode);
                        leftRotate(newNode);
                    }
                    setColor(parentOf(newNode),BLACK);
                    setColor(parentOf(parentOf(newNode)),RED);
                    rightRotate(parentOf(parentOf(newNode)));
                }
            }else{
                //看是不是4节点，有没有叔叔
                RBNode<K,V> y = leftOf(parentOf(parentOf(newNode)));
                if(colorOf(y) == RED){
                    //有叔叔
                    setColor(parentOf(parentOf(newNode)),RED);
                    setColor(parentOf(newNode),BLACK);
                    setColor(y,BLACK);
                    newNode = parentOf(parentOf(newNode));//递归到为空或是根节点
                }else{
                    //没有叔叔
                    if(newNode == leftOf(parentOf(newNode))){
                        newNode = parentOf(newNode);
                        rightRotate(newNode);
                    }
                    setColor(parentOf(newNode),BLACK);
                    setColor(parentOf(parentOf(newNode)),RED);
                    leftRotate(parentOf(parentOf(newNode)));
                }
            }
        }
        setColor(root,BLACK);
    }

    public RBNode<K,V> postNode(RBNode<K,V> node){
        if(node == null){
            return null;
        }else if(node.right != null){
            RBNode<K,V> p = rightOf(node);
            while(p.left != null){
                p = p.left;
            }
            return p;
        }else{
            RBNode<K,V> p = node.parent;
            RBNode<K,V> n = node;
            while(p != null && p.right == n){
                n = p;
                p = p.parent;
            }
            return p;
        }

    }

    public RBNode<K,V> getNode(K key){
        RBNode<K,V> node = root;
        while(node != null){
            int cmp = key.compareTo(node.key);
            if(cmp < 0){
                node = node.left;
            }else if(cmp > 0){
                node = node.right;
            }else{
                return node;
            }
        }
        return null;
    }

    public V remove(K key){
        RBNode<K,V> node = getNode(key);
        if(node == null){
            return null;
        }
        V oldValue = node.value;
        deleteNode(node);
        return oldValue;
    }

    public void deleteNode(RBNode<K,V> node){
        if(leftOf(node) != null && rightOf(node) != null){
            RBNode<K,V> postNode = postNode(node);
            node.key = postNode.key;
            node.value = postNode.value;
            node = postNode;
        }

        RBNode<K,V> replacement = node.left != null ? node.left : node.right;

        if(replacement != null){
            replacement.parent = node.parent;
            if(node.parent == null){
                root = replacement;
            }else if(node.parent.left == node){
                node.parent.left = replacement;
            }else{
                node.parent.right = replacement;
            }
            node.left = node.right = node.parent = null;
            if(node.color == BLACK){
                fixAfterRemove(replacement);
            }
        }else if(node.parent == null){
            root = null;
        }else{
            if(node.color == BLACK){
                fixAfterRemove(node);
            }
            if(node.parent != null){
                if(node.parent.left == node){
                    node.parent.left = null;
                }else{
                    node.parent.right = null;
                }
                node.parent = null;
            }
        }

    }

    public void fixAfterRemove(RBNode<K,V> node){
        while(node != root && colorOf(node) == BLACK){
            if(node == leftOf(parentOf(node))){
                RBNode<K,V> p = rightOf(parentOf(node));
                if(p.color == RED){
                    setColor(p,BLACK);
                    setColor(parentOf(node),RED);
                    leftRotate(parentOf(node));
                    p = rightOf(parentOf(node));
                }
                if(colorOf(leftOf(p)) == BLACK && colorOf(rightOf(p)) == BLACK){
                    //兄弟没有子节点，不借
                    setColor(p,RED);
                    p = p.parent;
                }else{
                    //兄弟有子节点，借
                    if(colorOf(rightOf(p)) == BLACK){
                        setColor(p,RED);
                        setColor(p.left,BLACK);
                        rightRotate(p);
                        p = rightOf(parentOf(node));
                    }
                    setColor(p,RED);
                    setColor(parentOf(p),BLACK);
                    setColor(rightOf(p),BLACK);
                    leftRotate(parentOf(node));
                    break;
                }
            }else{
                RBNode<K,V> p = leftOf(parentOf(node));
                if(p.color == RED){
                    //是不是真兄弟
                    setColor(p,BLACK);
                    setColor(parentOf(node),RED);
                    rightRotate(parentOf(node));
                    p = leftOf(parentOf(node));
                }
                if(colorOf(leftOf(p)) == BLACK && colorOf(rightOf(p)) == BLACK){
                    //兄弟没有子节点，不借
                    setColor(p,RED);
                    p = p.parent;
                }else{
                    //兄弟有子节点，借
                    if(colorOf(leftOf(p)) == BLACK){
                        setColor(p,RED);
                        setColor(p.right,BLACK);
                        leftRotate(p);
                        p = leftOf(parentOf(node));
                    }
                    setColor(p,RED);
                    setColor(parentOf(p),BLACK);
                    setColor(leftOf(p),BLACK);
                    rightRotate(parentOf(node));
                    break;
                }
            }
        }
        setColor(root,BLACK);
    }

    public void prePrint(){
        if(root == null){
            return;
        }else{
            root.prePrint();
        }
    }


    static class RBNode<K extends Comparable<K>,V>{
        private RBNode<K,V> parent;
        private RBNode<K,V> left;
        private RBNode<K,V> right;
        private K key;
        private V value;
        private Boolean color;

        public void prePrint(){
            System.out.println(this);
            if(left != null){
                left.prePrint();
            }
            if(right != null){
                right.prePrint();
            }
        }

        public RBNode(K key, V value, RBNode<K, V> parent) {
            this.key = key;
            this.value = value;
            this.parent = parent;
        }

        public RBNode(RBNode<K, V> parent, RBNode<K, V> left, RBNode<K, V> right, K key, V value, Boolean color) {
            this.parent = parent;
            this.left = left;
            this.right = right;
            this.key = key;
            this.value = value;
            this.color = color;
        }

        public RBNode() {
        }

        @Override
        public String toString() {
            return "RBNode{" +
                    "key=" + key +
                    ", value=" + value +
                    '}';
        }
    }


}
